package KnowSparkSQL.SQL.UDAF;

import org.apache.spark.sql.AnalysisException;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.expressions.MutableAggregationBuffer;
import org.apache.spark.sql.expressions.UserDefinedAggregateFunction;
import org.apache.spark.sql.types.*;

import java.util.Arrays;

/**
 * 下面我们自定义一个UDAF函数，
 * 名字起为mymax，它的功能和max函数一样，计算每组的某列数据的最大值。
 * 实现过程：
 * (1)自定义一个类，
 * 然后继承org.apache.spark.sql.expressions.UserDefinedAggregateFunction这个类。假设这个类的名字命名为MyMax
 * (2)父类UserDefinedAggregateFunction是一个抽象类，
 * MyMax类需要实现父类的抽象方法（即没有方法体的方法，每个方法的作用、输入参数、返回值等说明见代码注释）
 * (3)创建MyMax类型的对象，然后将该对象作为参数，传给UDF注册方法。
 */
public class UDAF_Demo {
    public static void main(String[] args) throws AnalysisException {
        SparkSession session = SparkSession.builder().appName(UDAF_Demo.class.getName()).master("local").getOrCreate();
        Dataset<Row> jsonStudent = session.read().json("./data/student.json");

        jsonStudent.createTempView("stu");

        session.udf().register("mymax", new MyMax());

        session.sql("select institute, sex, mymax(age) from stu group by institute,sex").show();




    }

    public static class MyMax extends UserDefinedAggregateFunction {

        @Override
        public StructType inputSchema() {
            StructType age1 = DataTypes.createStructType(Arrays.asList(new StructField("age1", DataTypes.LongType, true, Metadata.empty())));
            return age1;
        }

        @Override
        public StructType bufferSchema() {
            StructType age1 = DataTypes.createStructType(Arrays.asList(new StructField("age1", DataTypes.LongType, true, Metadata.empty())));
            return age1;
        }

        @Override
        public DataType dataType() {
            return DataTypes.LongType;
        }

        @Override
        public boolean deterministic() {
            return true;
        }

        @Override
        public void initialize(MutableAggregationBuffer mutableAggregationBuffer) {
            mutableAggregationBuffer.update(0, -1L);
        }

        @Override
        public void update(MutableAggregationBuffer mutableAggregationBuffer, Row row) {
            // 缓冲区
            long aLong = mutableAggregationBuffer.getLong(0);
            // 当前处理的行
            long aLong1 = row.getLong(0);
            if (aLong1 > aLong){
                mutableAggregationBuffer.update(0, aLong1);
            }


        }

        @Override
        public void merge(MutableAggregationBuffer mutableAggregationBuffer, Row row) {
            // 缓冲区
            long aLong = mutableAggregationBuffer.getLong(0);
            // 当前处理的行
            long aLong1 = row.getLong(0);
            if (aLong1 > aLong){
                mutableAggregationBuffer.update(0, aLong1);
            }



        }

        @Override
        public Object evaluate(Row row) {
            return row.getLong(0);
        }
    }
}
